package com.jq.jqaicode.ai.workflow;

import com.jq.jqaicode.ai.workflow.notes.*;
import com.jq.jqaicode.ai.workflow.ruotes.CodeTypeRoute;
import dev.langchain4j.data.message.ChatMessage;
import org.bsc.langgraph4j.CompiledGraph;
import org.bsc.langgraph4j.GraphStateException;
import org.bsc.langgraph4j.StateGraph;
import org.bsc.langgraph4j.langchain4j.serializer.std.LC4jStateSerializer;
import org.bsc.langgraph4j.prebuilt.MessagesState;
import org.bsc.langgraph4j.prebuilt.MessagesStateGraph;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.Map;

import static org.bsc.langgraph4j.StateGraph.END;
import static org.bsc.langgraph4j.StateGraph.START;
import static org.bsc.langgraph4j.action.AsyncNodeAction.node_async;
import static org.bsc.langgraph4j.action.AsyncEdgeAction.edge_async;

/**
 * 工作流部署
 */
@Configuration
public class WorkflowConfig {

    @Autowired
    HtmlCodeNode htmlCodeNode;
    @Autowired
    HtmlBuildNode htmlBuildNode;
    @Autowired
    VueCodeNode vueCodeNode;
    @Autowired
    VueBuildNode vueBuildNode;
    @Autowired
    CodeTypeRoute codeTypeRoute;


//                    start （用户输入）
//                      ↓
//                 codeTypeRoute（代码类型vue/html）
//                 ↙             ↘
//      htmlCode（代码生成）           vueCode（代码生成）
//           ↓                         ↓
//      htmlBuild（保存）           vueBuild（保存和部署）
//            ↘                      ↙
//             ↘                ↙
//                    end

    @Bean
    public CompiledGraph<MessagesState<ChatMessage>> messagesStateGraph() throws GraphStateException {
        // 构建workflow对象
        var stateSerializer = new LC4jStateSerializer<MessagesState<ChatMessage>>(MessagesState::new);
        StateGraph<MessagesState<ChatMessage>> stateGraph = new MessagesStateGraph<ChatMessage>(stateSerializer)
                .addNode("htmlCodeNode", node_async(htmlCodeNode))
                .addNode("vueCodeNode", node_async(vueCodeNode))
                .addNode("htmlBuildNode", node_async(htmlBuildNode))
                .addNode("vueBuildNode", node_async(vueBuildNode))

                .addConditionalEdges(START, edge_async(codeTypeRoute),
                        Map.of("html", "htmlCodeNode","vue", "vueCodeNode"))
                .addEdge("htmlCodeNode", "htmlBuildNode")
                .addEdge("htmlBuildNode", END)

                .addEdge("vueCodeNode", "vueBuildNode")
                .addEdge("vueBuildNode", END);
        return stateGraph.compile();
    }

}
